from scipy.misc import derivative
from scipy import integrate as sp_integrate
from sympy import Symbol, diff, cos, integrate
# 非windows
from jax import grad
from jax import numpy as jnp


class Derivative(object):
    """
    求导计算
    """
    def __init__(self, x_value):
        self.f = cos
        self.x_value = x_value

    def sympy_com(self):
        """
        sympy计算
        1、可以返回函数求导后表达式
        2、可多次求导

        :return: 对应值的导数
        """
        x = Symbol('x')
        # 一阶求导，返回函数表达式
        r = diff(self.f(x), x)
        # 二阶求导
        # r2 = diff(r, x)
        # 输入指定函数值，返回求导值
        return r.evalf(subs={'x': self.x_value})

    def sp_com(self, dx=1e-6):
        """
        scipy计算

        :return: 对应值的导数
        """
        return derivative(self.f, self.x_value, dx=dx)

    def jax_grad(self):
        """
        jax grad计算，暂不支持windows平台

        :return: 对应值的导数
        """
        grad_f = grad(jnp.cos)
        derivative_value = grad_f(self.x_value)
        return derivative_value


class Integral(object):
    """
    积分计算
    """
    def __init__(self):
        self.x = Symbol('x')
        self.f = 2 * self.x

    def sym_signal(self, start=0, end=1):
        """
        单层定积分: 对f在指定区间上求积分
        """
        # 返回积分表达式
        r = integrate(self.f, self.x)
        print('2 * x 的积分表达式： {}'.format(r))
        # 求积分值
        return integrate(self.f, (self.x, start, end))

    def sym_mul(self, out_start=0, out_end=3, in_start=0):
        """
        多次积分计算：内积分：[in_start, x]；外积分：[out_start, out_end]
        """
        t = Symbol('t')
        f = 2 * t
        # 内积分
        g = integrate(f, (t, in_start, self.x))
        # 外积分
        return integrate(g, (self.x, out_start, out_end))

    def sym_re(self, out_start=1, out_end=2, in_start=1):
        """
        二重积分：f(x, y) = x * y
        对y积分的范围 [1, x]
        对x积分的范围：[1, 2]
        """
        y = Symbol('y')
        f = self.x * y
        r = integrate(f, (self.x, out_start, out_end), (y, in_start, self.x))
        print('sym r 表达式: {}'.format(r))
        print('sym r 表达式参数为 2 的值: {}'.format(r.evalf(subs={'x': 2})))

    def sym_indefinite(self):
        return integrate(self.f, self.x)

    @staticmethod
    def sp_signal(start=0, end=1):
        """

        """
        def f(x):
            return 2 * x
        v, err = sp_integrate.quad(f, start, end)
        return v

    @staticmethod
    def sp_re(out_start=1, out_end=2, in_start=1):
        """
        二重积分：f(x, y) = x * y
        对y积分的范围 [1, x]
        对x积分的范围：[1, 2]
        """
        def f(x, y):
            return x * y

        def h(x):
            return x

        v, err = sp_integrate.dblquad(f, out_start, out_end,
                                      lambda x: in_start, h)
        return v


def run():
    """

    """
    # 求导
    x = 0.5
    der = Derivative(x_value=x)
    sym_value = der.sympy_com()
    sp_value = der.sp_com()
    jax_value = der.jax_grad()

    print('sympy计算结果： {}\nscipy计算结果：{}\njax grad计算结果：{}'.
          format(sym_value, sp_value, jax_value))

    # print('sympy计算结果： {}\nscipy计算结果：{}'.format(sym_value, sp_value))

    # 求积分
    inte = Integral()
    sigmal = inte.sym_signal()
    mul = inte.sym_mul()

    sp_sinal_v = inte.sp_signal()
    # 二重积分
    sp_re_v = inte.sp_re()

    info = '单层积分： sym:{},  sp: {}\n多次积分 sym:{}\n' \
           '不定积分 sym： {}\n二重积分 sp: {}'.\
        format(sigmal, sp_sinal_v, mul, inte.sym_indefinite(), sp_re_v)
    print(info)
    print('sym的二重积分:')
    inte.sym_re()


if __name__ == '__main__':
    run()
